// main.cpp
#include <vulkan/vulkan.h>
#include <vulkan/vulkan.hpp>
#include <vulkan/vulkan_core.h>
#include <iostream>
#include <vector>
#include <fstream>
#include <cstring>
#include <chrono>

using namespace std;
using namespace vk;
const int MATRIX_SIZE = 256;
#define VK_LOG_DEBUG(msg) std::cerr << msg << std::endl

// 矩阵数据结构
struct Matrix {
    std::vector<float> data;
    VkBuffer buffer;
    VkDeviceMemory memory;
};

// 初始化Vulkan实例和设备
static void initVulkan(VkInstance& instance, VkDevice& device, VkPhysicalDevice& physicalDevice, VkQueue& queue) {
    VK_LOG_DEBUG("enter initVulkan\n");
    // 创建实例
    VkApplicationInfo appInfo{};
    appInfo.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
    appInfo.pApplicationName = "VulkanMatMul";
    appInfo.applicationVersion = VK_MAKE_VERSION(1, 0, 0);
    appInfo.pEngineName = "No Engine";
    appInfo.engineVersion = VK_MAKE_VERSION(1, 0, 0);
    appInfo.apiVersion = VK_API_VERSION_1_0;

    // VK_LOG_DEBUG("before createInfo\n");
    VkInstanceCreateInfo createInfo{};
    createInfo.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
    createInfo.pApplicationInfo = &appInfo;
    const char* layers[] = {"VK_LAYER_KHRONOS_validation"};
    createInfo.enabledLayerCount = 1;
    createInfo.ppEnabledLayerNames = layers;

    
    vkCreateInstance(&createInfo, nullptr, &instance);
    
    // 选择物理设备
    uint32_t deviceCount = 0;
    vkEnumeratePhysicalDevices(instance, &deviceCount, nullptr);
    VK_LOG_DEBUG("deviceCount:" << deviceCount <<"\n");


    std::vector<VkPhysicalDevice> devices(deviceCount);
    vkEnumeratePhysicalDevices(instance, &deviceCount, devices.data());
    physicalDevice = devices[0]; // 默认选第一个

    VK_LOG_DEBUG("before get device information\n");
    

    // 3. 打印每个设备的信息
    for (const auto& device : devices) {
        // 设备属性
        VkPhysicalDeviceProperties deviceProperties;
        vkGetPhysicalDeviceProperties(device, &deviceProperties);
        
        // 设备特性
        VkPhysicalDeviceFeatures deviceFeatures;
        vkGetPhysicalDeviceFeatures(device, &deviceFeatures);
        
        // 设备内存属性
        VkPhysicalDeviceMemoryProperties memoryProperties;
        vkGetPhysicalDeviceMemoryProperties(device, &memoryProperties);

        std::cout << "Device Name: " << deviceProperties.deviceName << "\n";
        std::cout << "API Version: " << VK_VERSION_MAJOR(deviceProperties.apiVersion) << "."
                  << VK_VERSION_MINOR(deviceProperties.apiVersion) << "."
                  << VK_VERSION_PATCH(deviceProperties.apiVersion) << "\n";
        std::cout << "Driver Version: " << deviceProperties.driverVersion << "\n";
        std::cout << "Vendor ID: " << deviceProperties.vendorID << "\n";
        std::cout << "Device ID: " << deviceProperties.deviceID << "\n";
        std::cout << "Device Type: ";

        switch (deviceProperties.deviceType) {
            case VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU:
                std::cout << "Integrated GPU";
                break;
            case VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU:
                std::cout << "Discrete GPU";
                break;
            case VK_PHYSICAL_DEVICE_TYPE_VIRTUAL_GPU:
                std::cout << "Virtual GPU";
                break;
            case VK_PHYSICAL_DEVICE_TYPE_CPU:
                std::cout << "CPU";
                break;
            default:
                std::cout << "Other";
        }
        std::cout << "\n";

        // 打印内存堆信息
        std::cout << "\nMemory Heaps (" << memoryProperties.memoryHeapCount << "):\n";
        for (uint32_t i = 0; i < memoryProperties.memoryHeapCount; ++i) {
            std::cout << "  Heap " << i << ": Size = " << memoryProperties.memoryHeaps[i].size / (1024 * 1024) << " MB";
            if (memoryProperties.memoryHeaps[i].flags & VK_MEMORY_HEAP_DEVICE_LOCAL_BIT) {
                std::cout << " (Device Local)";
            }
            std::cout << "\n";
        }

        // 打印队列家族信息
        uint32_t queueFamilyCount = 0;
        vkGetPhysicalDeviceQueueFamilyProperties(device, &queueFamilyCount, nullptr);
        std::vector<VkQueueFamilyProperties> queueFamilies(queueFamilyCount);
        vkGetPhysicalDeviceQueueFamilyProperties(device, &queueFamilyCount, queueFamilies.data());

        std::cout << "\nQueue Families (" << queueFamilyCount << "):\n";
        for (uint32_t i = 0; i < queueFamilyCount; ++i) {
            std::cout << "  Family " << i << ": Count = " << queueFamilies[i].queueCount << ", Flags = ";
            if (queueFamilies[i].queueFlags & VK_QUEUE_GRAPHICS_BIT) std::cout << "Graphics ";
            if (queueFamilies[i].queueFlags & VK_QUEUE_COMPUTE_BIT) std::cout << "Compute ";
            if (queueFamilies[i].queueFlags & VK_QUEUE_TRANSFER_BIT) std::cout << "Transfer ";
            if (queueFamilies[i].queueFlags & VK_QUEUE_SPARSE_BINDING_BIT) std::cout << "Sparse ";
            std::cout << "\n";
        }

        // 打印设备扩展信息
        uint32_t extensionCount;
        vkEnumerateDeviceExtensionProperties(device, nullptr, &extensionCount, nullptr);
        std::vector<VkExtensionProperties> extensions(extensionCount);
        vkEnumerateDeviceExtensionProperties(device, nullptr, &extensionCount, extensions.data());

        std::cout << "\nDevice Extensions (" << extensionCount << "):\n";
        for (const auto& extension : extensions) {
            std::cout << "  " << extension.extensionName << " (Version " << extension.specVersion << ")\n";
        }

        std::cout << "\n----------------------------------------\n\n";


    }


    // 创建逻辑设备（启用计算队列）
    uint32_t queueFamilyIndex = 0;
    vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyIndex, nullptr);
    std::vector<VkQueueFamilyProperties> queueFamilies(queueFamilyIndex);
    vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyIndex, queueFamilies.data());

    for (uint32_t i = 0; i < queueFamilyIndex; i++) {
        if (queueFamilies[i].queueFlags & VK_QUEUE_COMPUTE_BIT) {
            queueFamilyIndex = i;
            break;
        }
    }

    float queuePriority = 1.0f;
    VkDeviceQueueCreateInfo queueCreateInfo{};
    queueCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
    queueCreateInfo.queueFamilyIndex = queueFamilyIndex;
    queueCreateInfo.queueCount = 1;
    queueCreateInfo.pQueuePriorities = &queuePriority;

    VkDeviceCreateInfo deviceCreateInfo{};
    deviceCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
    deviceCreateInfo.queueCreateInfoCount = 1;
    deviceCreateInfo.pQueueCreateInfos = &queueCreateInfo;

    try {
        vkCreateDevice(physicalDevice, &deviceCreateInfo, nullptr, &device);
    } catch (const vk::SystemError& e) {
        std::cerr<<e.what()<<std::endl;
    }
   

    VK_LOG_DEBUG("after vkGetDeviceQueue" << queue <<"\n");
    try {
        vkGetDeviceQueue(device, queueFamilyIndex, 0, &queue);
    }  catch (const vk::SystemError& e) {
        std::cerr<<e.what()<<std::endl;
    }
    VK_LOG_DEBUG("after vkGetDeviceQueue\n");
}

// 创建存储缓冲区
static void createBuffer(VkDevice device, VkPhysicalDevice physicalDevice, Matrix& matrix, VkBufferUsageFlags usage) {

    VkBufferCreateInfo bufferInfo{};
    bufferInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
    bufferInfo.size = matrix.data.size() * sizeof(float);
    bufferInfo.usage = usage;
    bufferInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
    vkCreateBuffer(device, &bufferInfo, nullptr, &matrix.buffer);

    VkMemoryRequirements memRequirements;
    vkGetBufferMemoryRequirements(device, matrix.buffer, &memRequirements);

    VkMemoryAllocateInfo allocInfo{};
    allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
    allocInfo.allocationSize = memRequirements.size;

    // 动态选择内存类型
    VkPhysicalDeviceMemoryProperties memProperties;
    vkGetPhysicalDeviceMemoryProperties(physicalDevice, &memProperties);
    for (uint32_t i = 0; i < memProperties.memoryTypeCount; i++) {
        if ((memRequirements.memoryTypeBits & (1 << i)) &&
            (memProperties.memoryTypes[i].propertyFlags & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT) &&
            (memProperties.memoryTypes[i].propertyFlags & VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)) {
            allocInfo.memoryTypeIndex = i;
            break;
        }
    }

    vkAllocateMemory(device, &allocInfo, nullptr, &matrix.memory);
    vkBindBufferMemory(device, matrix.buffer, matrix.memory, 0);

    // 复制数据到GPU
    void* data;
    vkMapMemory(device, matrix.memory, 0, bufferInfo.size, 0, &data);
    memcpy(data, matrix.data.data(), bufferInfo.size);
    vkUnmapMemory(device, matrix.memory);
}

// 创建Fence的正确流程
static VkFence createFence(VkDevice device) {
    VkFenceCreateInfo fenceInfo{};
    fenceInfo.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;
    // 默认初始状态为未触发（unsignaled）

    VkFence fence;
    if (vkCreateFence(device, &fenceInfo, nullptr, &fence) != VK_SUCCESS) {
        throw std::runtime_error("Failed to create fence!");
    }
    return fence;
}


int main() {
    


    VK_LOG_DEBUG("Hello Vulkan!\n");
    
    VkInstance instance;
    VkDevice device;
    VkPhysicalDevice physicalDevice;
    VkQueue queue;
    initVulkan(instance, device, physicalDevice, queue);
    VK_LOG_DEBUG("Vulkan initialized.\n");

    // 初始化矩阵
    Matrix a, b, c;
    a.data = std::vector<float>(MATRIX_SIZE * MATRIX_SIZE, 1.0f);
    b.data = std::vector<float>(MATRIX_SIZE * MATRIX_SIZE, 2.0f);
    c.data = std::vector<float>(MATRIX_SIZE * MATRIX_SIZE, 0.0f);

    // 创建缓冲区
    createBuffer(device, physicalDevice, a, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT);
    createBuffer(device, physicalDevice, b, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT);
    createBuffer(device, physicalDevice, c, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT);

    VK_LOG_DEBUG("Matrix initialized.");

    // 加载计算着色器（需预先编译为SPV）
    VkShaderModule shaderModule;
    std::ifstream file("matrix.comp.spv", std::ios::binary);
    std::vector<char> code((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
    VkShaderModuleCreateInfo shaderInfo{};
    shaderInfo.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
    shaderInfo.codeSize = code.size();
    shaderInfo.pCode = reinterpret_cast<const uint32_t*>(code.data());
    vkCreateShaderModule(device, &shaderInfo, nullptr, &shaderModule);

    VK_LOG_DEBUG("Shader loaded.");

    // 创建描述符集布局
    VkDescriptorSetLayoutBinding bindings[3];
    for (int i = 0; i < 3; i++) {
        bindings[i].binding = i;
        bindings[i].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
        bindings[i].descriptorCount = 1;
        bindings[i].stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;
        bindings[i].pImmutableSamplers = nullptr;
    }

    VK_LOG_DEBUG("Descriptor set layout created.");

    VkDescriptorSetLayoutCreateInfo layoutInfo{};
    layoutInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
    layoutInfo.bindingCount = 3;
    layoutInfo.pBindings = bindings;
    VkDescriptorSetLayout descriptorSetLayout;
    vkCreateDescriptorSetLayout(device, &layoutInfo, nullptr, &descriptorSetLayout);

    VK_LOG_DEBUG("Descriptor set layout created.");

    

    // 创建计算管线
    VkPipelineLayoutCreateInfo pipelineLayoutInfo{};
    pipelineLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
    pipelineLayoutInfo.setLayoutCount = 1;
    pipelineLayoutInfo.pSetLayouts = &descriptorSetLayout;
    VkPipelineLayout pipelineLayout;
    vkCreatePipelineLayout(device, &pipelineLayoutInfo, nullptr, &pipelineLayout);

    VK_LOG_DEBUG("Pipeline layout created.");

    VkComputePipelineCreateInfo pipelineInfo{};
    pipelineInfo.sType = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO;
    pipelineInfo.stage.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    pipelineInfo.stage.stage = VK_SHADER_STAGE_COMPUTE_BIT;
    pipelineInfo.stage.module = shaderModule;
    pipelineInfo.stage.pName = "main";
    pipelineInfo.layout = pipelineLayout;
    VkPipeline pipeline;
    vkCreateComputePipelines(device, VK_NULL_HANDLE, 1, &pipelineInfo, nullptr, &pipeline);

    VK_LOG_DEBUG("Pipeline created.");

    // 创建描述符池和描述符集
    VkDescriptorPoolSize poolSize{};
    poolSize.type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
    poolSize.descriptorCount = 3;

    VkDescriptorPoolCreateInfo poolInfo{};
    poolInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
    poolInfo.poolSizeCount = 1;
    poolInfo.pPoolSizes = &poolSize;
    poolInfo.maxSets = 1;
    VkDescriptorPool descriptorPool;
    vkCreateDescriptorPool(device, &poolInfo, nullptr, &descriptorPool);

    VkDescriptorSetAllocateInfo allocInfo{};
    allocInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
    allocInfo.descriptorPool = descriptorPool;
    allocInfo.descriptorSetCount = 1;
    allocInfo.pSetLayouts = &descriptorSetLayout;
    VkDescriptorSet descriptorSet;
    vkAllocateDescriptorSets(device, &allocInfo, &descriptorSet);

   

    VkDescriptorBufferInfo bufferInfos[3] = {
        {a.buffer, 0, VK_WHOLE_SIZE},
        {b.buffer, 0, VK_WHOLE_SIZE},
        {c.buffer, 0, VK_WHOLE_SIZE}
    };

    VkWriteDescriptorSet writes[3];
    memset(writes, 0, sizeof(writes)); // 初始化为
    for (int i = 0; i < 3; i++) {
        writes[i].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
        writes[i].pNext = nullptr; // 显式设置pNext
        writes[i].dstSet = descriptorSet;
        writes[i].dstBinding = i;
        writes[i].descriptorCount = 1;
        writes[i].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
        writes[i].pBufferInfo = &bufferInfos[i];
    }
    
    vkUpdateDescriptorSets(device, 3, writes, 0, nullptr);
    
    VK_LOG_DEBUG("Descriptor set created.");
    
    // 录制命令缓冲区
    VkCommandPool commandPool;
    VkCommandPoolCreateInfo poolCreateInfo{};
    poolCreateInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
    poolCreateInfo.queueFamilyIndex = 0; // 需与队列族一致
    vkCreateCommandPool(device, &poolCreateInfo, nullptr, &commandPool);

    VkCommandBuffer commandBuffer;
    VkCommandBufferAllocateInfo allocCmdBufInfo{};
    allocCmdBufInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
    allocCmdBufInfo.commandPool = commandPool;
    allocCmdBufInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    allocCmdBufInfo.commandBufferCount = 1;
    vkAllocateCommandBuffers(device, &allocCmdBufInfo, &commandBuffer);

    VkCommandBufferBeginInfo beginInfo{};
    beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
    vkBeginCommandBuffer(commandBuffer, &beginInfo);

    vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline);
    vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, pipelineLayout, 0, 1, &descriptorSet, 0, nullptr);
    vkCmdDispatch(commandBuffer, MATRIX_SIZE/16, MATRIX_SIZE/16, 1); // 工作组大小16x16 [[6]]

    vkEndCommandBuffer(commandBuffer);

    // 提交命令并同步
    VkSubmitInfo submitInfo{};
    submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
    submitInfo.commandBufferCount = 1;
    submitInfo.pCommandBuffers = &commandBuffer;

        // 在命令提交时使用
    VkFence fence = createFence(device);

    auto start = std::chrono::high_resolution_clock::now();
    vkQueueSubmit(queue, 1, &submitInfo, fence); // 提交时关联Fence
    vkWaitForFences(device, 1, &fence, VK_TRUE, UINT64_MAX);
    auto end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> elapsed = end - start;
    std::cout << "Total time: " << elapsed.count() * 1000 << " ms" << std::endl;

    // 读取结果
    void* mappedMemory;
    vkMapMemory(device, c.memory, 0, VK_WHOLE_SIZE, 0, &mappedMemory);
    std::vector<float> result(MATRIX_SIZE * MATRIX_SIZE);
    memcpy(result.data(), mappedMemory, MATRIX_SIZE * MATRIX_SIZE * sizeof(float));
    vkUnmapMemory(device, c.memory);

    // 验证结果（预期每个元素为256*2=512）
    std::cout << "Result[0][0] = " << result[0] << std::endl;
    // 在memcpy获取结果后添加以下代码：
    std::cout << "Matrix Result:\n";
    // for (int row = 0; row < MATRIX_SIZE; ++row) {
    //     for (int col = 0; col < MATRIX_SIZE; ++col) {
    //         std::cout << result[row * MATRIX_SIZE + col] << "\t";
    //     }
    //     std::cout << "\n";
    // }

    // 清理资源（省略部分代码）
    vkDestroyFence(device, fence, nullptr);
    vkDestroyCommandPool(device, commandPool, nullptr);
    vkDestroyPipeline(device, pipeline, nullptr);
    vkDestroyPipelineLayout(device, pipelineLayout, nullptr);
    vkDestroyDescriptorPool(device, descriptorPool, nullptr);
    vkDestroyDescriptorSetLayout(device, descriptorSetLayout, nullptr);
    vkDestroyShaderModule(device, shaderModule, nullptr);
    vkDestroyBuffer(device, a.buffer, nullptr);
    vkFreeMemory(device, a.memory, nullptr);
    vkDestroyBuffer(device, b.buffer, nullptr);
    vkFreeMemory(device, b.memory, nullptr);
    vkDestroyBuffer(device, c.buffer, nullptr);
    vkFreeMemory(device, c.memory, nullptr);
    vkDestroyDevice(device, nullptr);
    vkDestroyInstance(instance, nullptr);

    return 0;
}